Skip to content

feat(npu): capture AsyncCAM layered FFN in a reusable graph - #418

Open
jiangkuaixue123 wants to merge 8 commits into
vllm-project:mainfrom
jiangkuaixue123:codex/cam-async-ffn-graph
Open

jiangkuaixue123 wants to merge 8 commits into
vllm-project:mainfrom
jiangkuaixue123:codex/cam-async-ffn-graph

Conversation

@jiangkuaixue123

@jiangkuaixue123 jiangkuaixue123 commented Sep 30, 2026 •

Copy link
Copy Markdown
Collaborator

Purpose

Capture a reusable AsyncCAM layered FFN graph for DeepSeek V4 W4A8, while keeping cross-role startup coordination outside the connector and worker. Eager FFN remains supported; non-eager FFN requires FULL. Attention remains eager.

Scope

One fixed-capacity dispatch_recv → layered GMM → combine_send graph per FFN rank, startup coordination, and focused runtime tests. Native kernels, native ABI, graph capacity, normal data wire format, and upstream checkouts are unchanged.

Implementation Notes

  • A typed ProcessGroupRendezvousContext borrows the existing rendezvous Store and process group. Factory injection is optional and limited to CAM; existing connector creation calls retain their behavior.
  • AsyncCamStartupCoordinator owns nonce isolation, mode agreement, paired warmup, readiness, and failure propagation. The connector owns communication resources; runners own computation/graphs; the worker owns its receive thread.
  • Prepare the layered executor in both eager and FULL modes without standalone GMM execution. FULL startup warms up through existing connector send/receive interfaces, consuming the Attention FIFO before profile. All ranks complete communication warmup before FFN capture; Attention waits for FFN capture and receive-thread readiness before returning from load_model.
  • Sticky failures prevent restarting a failed startup, and ready cannot overwrite failure. Startup Store polling does not enter the work-item hot path.
  • Close the graph and connector only after the receive thread confirms device completion. Blocked or failed receivers retain resources for process-level cleanup. No device cancellation or graceful graph-stop mechanism is added.

Test Plan

Focused tests cover coordinator ordering/failure propagation, real connector FIFO consumption, Factory compatibility, rendezvous context lifecycle, runner/worker shutdown, and native Meta contracts. Full-service validation uses jcz_afd1, DSV4-Flash prefill_only, Attention DP2×TP4 and FFN EP8, with seven single-token requests per FULL/eager deployment: short, medium, 25,506-token cross-chunk, and four concurrent requests.

Test Result

  • Frozen revision 4e86c71: all 16 PR source files on jcz_afd1 matched committed file hashes. Actual module import paths pointed to the synced checkout; native library hashes remained those of the original kernels. Process commands, working directories, and start times were recorded for both deployments.
  • All 152 focused tests passed. After adding two startup-only diagnostic messages, all nine coordinator cases passed again. Pinned Ruff 0.15.13 check/format, Python compilation and diff checks passed.
  • FULL: Attention standard-interface DS/CR warmup completed on all eight ranks; each FFN rank consumed two work items for Attention DP2, captured exactly one graph, and published graph-mode ready. All Attention ranks observed FFN readiness before engine initialization completed. All seven API requests passed.
  • Eager: all eight FFN ranks published eager-mode ready, no graphs were captured, and all seven API requests passed. Single-token outputs matched FULL case by case, including the 25,506-token cross-chunk prompt and four concurrent requests.
  • Tested vLLM 2e882b2b64ac, vLLM-Ascend 508697610699, CANN 9.0.1 and torch_npu 2.10.0.post2. Validation used process-level cleanup; graceful shutdown, active-request cancellation, full generation accuracy, and performance improvements are not claimed.

Docs Impact

No Markdown files included, as requested. Local plans and validation reports remain outside this PR.

Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
@jiangkuaixue123
jiangkuaixue123 force-pushed the codex/cam-async-ffn-graph branch from 092b5c4 to 57d145c Compare September 30, 2026 07:19
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
@hsliuustc0106 hsliuustc0106 added enhancement New feature or request Ascend Ascend NPU platform and related changes labels Sep 30, 2026
jiangkuaixue123 added a commit that referenced this pull request Oct 10, 2026
…sable graph (#436)

## Purpose

Backport #418 to `release/v0.26.0rc2`: capture a reusable AsyncCAM
layered FFN graph for DeepSeek V4 W4A8, with cross-role startup
coordination. Eager FFN remains supported; non-eager FFN requires FULL.
Attention remains eager.

## Changes

Cherry-picked all eight commits from #418, through
`4e86c7182e570c650f98d6c6b45910fe084f99eb`, without conflicts or
additional adaptations. All 16 changed files are byte-for-byte identical
to the source PR head.

The startup coordinator owns mode agreement, paired warmup, readiness,
and failure propagation. Runners own computation and graphs, the
connector owns communication resources, and the worker owns its receive
thread. Native kernels and ABI remain unchanged.

## Validation

- Ruff 0.15.13 lint and format checks passed for all 16 changed files.
- Python compilation and `git diff --check` passed.
- Attempted the three focused unit-test modules: all three were skipped
because the local environment lacks `torch` / `torch_npu`; no runtime
tests executed locally.
- NPU service validation was not rerun for this release branch. See #418
for the original FULL/eager validation results; those results are
source-PR evidence, not validation of this backport.

---------

Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
jiangkuaixue123 pushed a commit that referenced this pull request Oct 12, 2026
…h Store (#448)

## Purpose

Fix AsyncCAM startup with different Attention and FFN HCCL buffer sizes.
With Attention at 1024 MiB and FFN at 4096 MiB, `_get_store()` fails in
`HcomBroadcast` with inconsistent `inCclBufferSize` / `outCclBufferSize`
metadata. This path runs in eager mode as well as FULL mode.

Exchange the startup namespace through the existing rendezvous Store so
this host metadata no longer requires an HCCL collective.

## Issue

Related to #436, which backported the startup coordinator from #418.
Based on `release/v0.26.0rc2` at `cda849a5`.

## Scope

- In scope: `_get_store()` and focused startup tests; production diff is
+8/-20 lines.
- Out of scope: CAM data transfer, native kernels, inference scheduling,
and recovery using a surviving Store after partial rank restarts.

## Implementation Notes

- Rank 0 writes a random namespace to a bootstrap `PrefixStore`; all
ranks wait up to the existing 300-second startup timeout, read it, and
cache the resulting namespace-prefixed Store.
- This relies on one startup per fresh CAM rendezvous Store.
Cached-context validation and subsequent mode/ready/failure coordination
are preserved.
- Remove the nonce tensor, broadcast, device synchronization, and
device-to-host readback from `_get_store()`. Existing
communication-warmup synchronizations remain in place; Store readiness
is not a device-completion fence.

## Test Plan

- Run `python -m pytest -q --confcutdir=tests
tests/unit/v1/worker/test_async_cam_startup.py` in a torch_npu
environment. New cases use a real HashStore with 16 threads to cover
leader-first/leader-late arrival, cached access, isolation across fresh
Stores, and missing-leader timeout; device nonce operations are
forbidden.
- Run Ruff check/format and `git diff --check` for the two changed
files.
- On the same 16 DIE, compare the previous local Store handshake
workaround with this patch, using 12A+4F, eager mode, A1G/F4G, HTTP
concurrency 16, and identical frozen inputs/routes. Each cold-cache run
contains 16 complete sessions / 328 requests, with one output token per
request.

## Test Result

- Startup unit tests: **12 passed, 0 skipped**. Ruff 0.15.13
check/format and diff check passed.
- Both hardware runs: **328/328 successful**, with matching startup
arguments and performance environment variables. Only this startup
module differs between the deployed packages.
- Cached tokens and actual prefill tokens match exactly: 18,604,032 and
1,904,633 respectively, with no preemptions. Request routing, dependency
order, and concurrency audits passed.

| Metric | Previous Store workaround | This patch | Change |
|---|---:|---:|---:|
| Replay duration | 104.94 s | 105.31 s | +0.36% |
| Throughput | 3.126 requests/s | 3.115 requests/s | -0.36% |
| TTFT P95 | 7.115 s | 7.090 s | -0.36% |

The unpatched HCCL-broadcast version cannot start with this A1G/F4G
configuration, so it is not the performance baseline. This single paired
run shows essentially unchanged performance, not statistical proof of
zero regression. Hardware coverage is limited to eager A1G/F4G; FULL
mode and arbitrary buffer sizes were not revalidated. Natural cache
allocation differed by three blocks in one domain (0.0053%); measured
cache/prefill work was identical.

## Docs Impact

No documentation files changed. The fresh-Store lifecycle assumption is
documented next to the implementation; no configuration option or public
API is added.

---

<details>
<summary>Essential PR Checklist</summary>

- [x] Purpose is clear and linked to public context.
- [x] Scope is bounded.
- [x] Compatibility with vLLM v0.26.0 is considered.
- [x] No changes are made to the vLLM source checkout.
- [x] The change stays in the plugin-owned startup coordinator.
- [x] No compatibility shim or runtime monkey patch is added.
- [x] No new device-dependent import or import-time work is added.
- [x] Validation evidence and hardware coverage limits are included.
- [x] Documentation impact is stated.

</details>

Signed-off-by: jiaran-king <jiaranran2@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Ascend Ascend NPU platform and related changes enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants